import subprocess
import time
import os

PREFIX = os.path.dirname(__file__)


class MixedDeploymentTester:
    def __init__(self):
        self.wrk_threads = 80
        self.wrk_connections = 100
        self.test_duration = 120  # seconds
        self.stress_cpu_load = 100
        self.stress_vm_bytes = "90%"
        self.stress_run_time = "180s"
        self.nginx_namespace = "default"
        self.latency_data = []

    def deploy_nginx(self):
        # 部署Nginx在线服务
        nginx_deployment_path = os.path.join(PREFIX, 'nginx_deployment.yaml')
        nginx_service_path = os.path.join(PREFIX, 'nginx_service.yaml')
        cmd_deployment = [
            "kubectl",
            "apply",
            "-f", nginx_deployment_path,
            "-n", self.nginx_namespace
        ]
        subprocess.run(cmd_deployment, check=True)

        cmd_service = [
            "kubectl",
            "apply",
            "-f", nginx_service_path,
            "-n", self.nginx_namespace
        ]
        subprocess.run(cmd_service, check=True)
        
        self._wait_for_nginx()

    def _wait_for_nginx(self):
        """等待Nginx就绪"""
        start_time = time.time()
        temp = '''
-o=jsonpath='{range .items[*]}{@.status.conditions[?(@.type=="Ready")].status}{"\\n"}{end}'
'''
        while True:
            cmd = [
                "kubectl",
                "get",
                "pods",
                "-n", self.nginx_namespace,
                "-l", "app=nginx",
                temp.strip()
            ]
            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            ready_pods = result.stdout.count('True')

            if ready_pods:
                break

            if time.time() - start_time > 60:
                raise RuntimeError("nginx启动超时")
            time.sleep(2) 

    def _get_nginx_service_url(self):
        """获取Nginx访问地址"""
        cmd_get_ip = [
            "kubectl",
            "get",
            "svc",
            "nginx-service",
            "-n", self.nginx_namespace,
            "-o", "jsonpath='{.spec.clusterIP}'"
        ]
        cmd_get_port = [
            "kubectl",
            "get",
            "svc",
            "nginx-service",
            "-n", self.nginx_namespace,
            "-o", "jsonpath='{.spec.ports[0].port}'"
        ]
        ip_result = subprocess.run(cmd_get_ip, capture_output=True, text=True)
        port_result = subprocess.run(cmd_get_port, capture_output=True, text=True)
        ip_result = ip_result.stdout.strip("'")
        port_result = port_result.stdout.strip("'")
        url = f"http://{ip_result}:{port_result}"   
        return url

    def run_wrk_test(self, scenario_name):
        """执行WRK性能测试"""
        url = self._get_nginx_service_url()        
        cmd = [
            "wrk",
            f"-t {self.wrk_threads}",
            f"-c {self.wrk_connections}",
            f"-d {self.test_duration}s",
            "--latency",
            url
        ]        
        result = subprocess.run(cmd, capture_output=True, text=True)
        # 测试结果：result.stdout

    def run_stress_test(self, stress_ng_pod_name):
        """启动离线压力测试"""
        cmd = [
            "kubectl",
            "exec", stress_ng_pod_name,
            "--",
            "stress-ng",
            "--cpu", "16",
            "--cpu-method", "pi",
            "--cpu-load", str(self.stress_cpu_load),
            "--vm", "1",
            "--vm-bytes", self.stress_vm_bytes,
            "--vm-keep",
            "-t", self.stress_run_time
        ]
        process = subprocess.Popen(cmd, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        time.sleep(45)
        return process
    
    @staticmethod
    def _check_stress_test():
        '''检查stress-ng离线任务容器是否完全启动'''
        cmd = [
            "kubectl",
            "top",
            "nodes"
        ]
        subprocess.run(cmd, check=True)

    def cleanup(self, stress_process):
        """清理测试资源"""
        cmd_del_deployment = [
            "kubectl",
            "delete",
            "deployments",
            "-l", "app=nginx-server-deployment",
            "-n", self.nginx_namespace
        ]
        cmd_del_service = [
            "kubectl",
            "delete",
            "services",
            "-l", "app=nginx-server-service",
            "-n", self.nginx_namespace
        ]
        subprocess.run(cmd_del_deployment, check=False)
        subprocess.run(cmd_del_service, check=False)
        if stress_process:
            stress_process.terminate()
            time.sleep(1)
            if stress_process.poll() is None:
                stress_process.kill()

    def pipeline(self, test_mode, stress_pod_name):
        stress_process = ''
        # 启动Nginx服务
        self.deploy_nginx()

        # 部署stress-ng离线服务容器
        if test_mode in "With Stress":
            stress_process = self.run_stress_test(stress_pod_name)
            self._check_stress_test()

        # 执行WRK性能测试
        self.run_wrk_test(test_mode)
        self._check_stress_test()

        # 清除测试资源
        self.cleanup(stress_process)


if __name__ == "__main__":
    tester = MixedDeploymentTester()

    # 启动进程访问延时变化测试：输入1为无压场景，2为有压场景
    test_mode_list = ['No Stress', 'With Stress']
    mode_input = input()
    if mode_input == '1':
        test_mode = test_mode_list[0]
    else:
        test_mode = test_mode_list[1]

    stress_pod_name = 'stress-ng-deployment-0620-67bd678847-nhblv'

    tester.pipeline(test_mode, stress_pod_name)